﻿using System.Collections;
using System.Collections.Generic;
using Unity.Mathematics;
using UnityEngine;
using UnityEngine.Rendering;

public class GPUTerrain : MonoBehaviour
{
    public Terrain terrain;
    public Mesh instanceMesh;
    public Material terrainMat;
    public ComputeShader cullingComputeShader;

    public int DebugMode = 0;

    private List<NodeInfo> allNodeInfo;
    private RenderTexture heightmapTex;
    private Texture2D normalTex;
    private uint[] computeArgs;
    private uint[] secondPassComputeArgs;

    private ComputeBuffer allInstancesPosWSBuffer;
    private ComputeBuffer visibleInstancesOnlyPosWSIDBuffer;
    private ComputeBuffer invisibleInstancesOnlyPosWSIDBuffer;
    private ComputeBuffer totalvisibleInstancesOnlyPosWSIDBuffer;
    private ComputeBuffer argsBuffer;
    private ComputeBuffer computeArgBuffer;
    private int cullTerrainKernel;
    private int cullTerrainKernelSecondPass;
    private TerrainNodePage pageRoot;
    const string m_ProfilerTag = "Gpu Terrain";
    const string m_ProfilerPrepassTag = "Gpu Terrain Depth Prepass";

    void Start()
    {
        if (pageRoot == null)
        {
            float perSize = 64;
            Rect rect = new Rect(0, 0, terrain.terrainData.size.x, terrain.terrainData.size.z);
            pageRoot = new TerrainNodePage(rect);
            var children = new List<TerrainNodePage>();
            for (var i = rect.xMin; i < rect.xMax; i += perSize)
                for (var j = rect.yMin; j < rect.yMax; j += perSize)
                {
                    children.Add(new TerrainNodePage(new Rect(i, j, perSize, perSize), 3));
                }
            pageRoot.children = children.ToArray();
        }
        allNodeInfo = new List<NodeInfo>();
        var camPos = Camera.main.transform.position;
        var center = new Vector2(camPos.x, camPos.z);
        pageRoot.CollectNodeInfo(center, allNodeInfo);
        for (int i = 0; i < allNodeInfo.Count; i++)
        {
            var nodeInfo = allNodeInfo[i];
            var nodeCenter = new Vector2(nodeInfo.rect.x + nodeInfo.rect.z * 0.5f, nodeInfo.rect.y + nodeInfo.rect.w * 0.5f);
            var topNode = pageRoot.GetActiveNode(nodeCenter + new Vector2(0, nodeInfo.rect.w));
            var bottomNode = pageRoot.GetActiveNode(nodeCenter + new Vector2(0, -nodeInfo.rect.w));
            var leftNode = pageRoot.GetActiveNode(nodeCenter + new Vector2(-nodeInfo.rect.z, 0));
            var rightNode = pageRoot.GetActiveNode(nodeCenter + new Vector2(nodeInfo.rect.z, 0));
            // 记录当前Node跟周围node的mip的差别，通过移动相差mip的中间点，来避免出现缝隙
            var nei = new bool4(topNode != null && topNode.mip > nodeInfo.mip,
                bottomNode != null && bottomNode.mip > nodeInfo.mip,
                leftNode != null && leftNode.mip > nodeInfo.mip,
                rightNode != null && rightNode.mip > nodeInfo.mip);
            nodeInfo.neighbor = (1 * (nei.x ? 1 : 0)) + ((1 << 1) * (nei.y ? 1 : 0)) + ((1 << 2) * (nei.z ? 1 : 0)) + ((1 << 3) * (nei.w ? 1 : 0));
            allNodeInfo[i] = nodeInfo;
        }

        heightmapTex = terrain.terrainData.heightmapTexture;
        normalTex = new Texture2D(heightmapTex.width, heightmapTex.height, TextureFormat.RGBA32, -1, true);
        var colors = new Color[heightmapTex.width * heightmapTex.width];
        int intdex = 0;
        for (int i = 0; i < heightmapTex.width; ++i)
        {
            for (int j = 0; j < heightmapTex.height; ++j)
            {
                var normal = terrain.terrainData.GetInterpolatedNormal((float)i / heightmapTex.width, (float)j / heightmapTex.height);
                colors[intdex++] = new Color(normal.z * 0.5f + 0.5f, normal.y * 0.5f + 0.5f, normal.x * 0.5f + 0.5f);
            }
        }
        normalTex.SetPixels(colors);
        normalTex.Apply();

        allInstancesPosWSBuffer = new ComputeBuffer(allNodeInfo.Count, sizeof(float) * 4 + sizeof(int) + sizeof(int));
        allInstancesPosWSBuffer.SetData(allNodeInfo.ToArray());

        visibleInstancesOnlyPosWSIDBuffer = new ComputeBuffer(allNodeInfo.Count, sizeof(uint), ComputeBufferType.Append);
        invisibleInstancesOnlyPosWSIDBuffer = new ComputeBuffer(allNodeInfo.Count, sizeof(uint), ComputeBufferType.Append);
        totalvisibleInstancesOnlyPosWSIDBuffer = new ComputeBuffer(allNodeInfo.Count, sizeof(uint), ComputeBufferType.Append);

        if (argsBuffer != null)
            argsBuffer.Release();
        uint[] args = new uint[5] { 0, 0, 0, 0, 0 };
        argsBuffer = new ComputeBuffer(1, args.Length * sizeof(uint), ComputeBufferType.IndirectArguments);

        args[0] = (uint)instanceMesh.GetIndexCount(0);
        args[1] = (uint)allNodeInfo.Count;
        args[2] = (uint)instanceMesh.GetIndexStart(0);
        args[3] = (uint)instanceMesh.GetBaseVertex(0);
        args[4] = 0;

        argsBuffer.SetData(args);

        if (computeArgBuffer != null)
            computeArgBuffer.Release();
        computeArgs = new uint[3] { 1, 1, 1 };
        computeArgBuffer = new ComputeBuffer(1, computeArgs.Length * sizeof(uint), ComputeBufferType.IndirectArguments);

        computeArgs[0] = (uint)Mathf.CeilToInt(allNodeInfo.Count / 64.0f);
        computeArgs[1] = 1;
        computeArgs[2] = 1;

        computeArgBuffer.SetData(computeArgs);

        secondPassComputeArgs = new uint[3] { 1, 1, 1 };

        cullTerrainKernel = this.cullingComputeShader.FindKernel("CullTerrain");
        cullingComputeShader.SetBuffer(cullTerrainKernel, "_AllInstancesPosWSBuffer", allInstancesPosWSBuffer);
        cullingComputeShader.SetTexture(cullTerrainKernel, "_HeightMap", heightmapTex);
        cullingComputeShader.SetFloat("_TerrainHeightSize", 2 * terrain.terrainData.size.y);

        cullTerrainKernelSecondPass = this.cullingComputeShader.FindKernel("CullTerrainInstance");
        cullingComputeShader.SetBuffer(cullTerrainKernelSecondPass, "_AllInstancesPosWSBuffer", allInstancesPosWSBuffer);
        cullingComputeShader.SetTexture(cullTerrainKernelSecondPass, "_HeightMap", heightmapTex);
        cullingComputeShader.SetFloat("_TerrainHeightSize", 2 * terrain.terrainData.size.y);

        terrainMat.SetBuffer("_AllInstancesTransformBuffer", allInstancesPosWSBuffer);
        terrainMat.SetBuffer("_VisibleInstanceOnlyTransformIDBuffer", totalvisibleInstancesOnlyPosWSIDBuffer);

        Shader.SetGlobalTexture("_TerrainHeightmapTexture", heightmapTex);
        Shader.SetGlobalTexture("_TerrainNormalmapTexture", normalTex);
        Shader.SetGlobalVector("terrainParam", terrain.terrainData.size);

        GPUTerrainPass.ExecuteAction += Render;
        TerrainDepthOnlyPass.ExecuteAction += DepthPreRender;
    }

    private void OnDestroy()
    {
        if (allInstancesPosWSBuffer != null)
            allInstancesPosWSBuffer.Release();
        allInstancesPosWSBuffer = null;

        if (visibleInstancesOnlyPosWSIDBuffer != null)
            visibleInstancesOnlyPosWSIDBuffer.Release();
        visibleInstancesOnlyPosWSIDBuffer = null;

        if (invisibleInstancesOnlyPosWSIDBuffer != null)
            invisibleInstancesOnlyPosWSIDBuffer.Release();
        invisibleInstancesOnlyPosWSIDBuffer = null;

        if (totalvisibleInstancesOnlyPosWSIDBuffer != null)
            totalvisibleInstancesOnlyPosWSIDBuffer.Release();
        totalvisibleInstancesOnlyPosWSIDBuffer = null;

        if (argsBuffer != null)
            argsBuffer.Release();
        argsBuffer = null;

        if (computeArgBuffer != null)
            computeArgBuffer.Release();
        computeArgBuffer = null;

        if (normalTex != null)
            DestroyImmediate(normalTex);
        normalTex = null;

        GPUTerrainPass.ExecuteAction -= Render;
        TerrainDepthOnlyPass.ExecuteAction -= DepthPreRender;
    }

    // https://gamedev.net/forums/topic/701850-accessing-appendstructuredbuffer-count/5404234/
    // https://forum.unity.com/threads/access-buffer-counter-in-compute-shader.656722/
    void Render(ScriptableRenderContext context, Camera cam)
    {
        var cmd = CommandBufferPool.Get(m_ProfilerTag);
        context.ExecuteCommandBuffer(cmd);
        cmd.Clear();

        // if (DebugMode < 0 || (DebugMode == 0 && cam == Camera.main))
        if (DebugMode < 0 || (DebugMode == 0))
        {
            cmd.CopyCounterValue(totalvisibleInstancesOnlyPosWSIDBuffer, argsBuffer, 4);
            cmd.DrawMeshInstancedIndirect(instanceMesh, 0, terrainMat, 0, argsBuffer);
        }

        context.ExecuteCommandBuffer(cmd);
        CommandBufferPool.Release(cmd);
    }

    void DepthPreRender(ScriptableRenderContext context, Camera cam, RenderTexture hizRT, int drawIndex, Matrix4x4 vpPrev)
    {
        var cmd = CommandBufferPool.Get(m_ProfilerPrepassTag);
        context.ExecuteCommandBuffer(cmd);
        cmd.Clear();



        if (DebugMode < 0 || (DebugMode == 0 && cam == Camera.main))
        {

            if (drawIndex == 0)
            {
                //computeArgBuffer.SetData(computeArgs);

                visibleInstancesOnlyPosWSIDBuffer.SetCounterValue(0);
                invisibleInstancesOnlyPosWSIDBuffer.SetCounterValue(0);
                totalvisibleInstancesOnlyPosWSIDBuffer.SetCounterValue(0);
                cmd.SetComputeTextureParam(cullingComputeShader, cullTerrainKernel, "_HiZMap", hizRT);
                cmd.SetComputeVectorParam(cullingComputeShader, "_HizSize", new Vector4(hizRT.width, hizRT.height, 0, 0));

                cmd.SetComputeBufferParam(cullingComputeShader, cullTerrainKernel, "_VisibleInstancesOnlyPosWSIDBuffer", visibleInstancesOnlyPosWSIDBuffer);
                cmd.SetComputeBufferParam(cullingComputeShader, cullTerrainKernel, "_InVisibleInstancesOnlyPosWSIDBuffer", invisibleInstancesOnlyPosWSIDBuffer);
                cmd.SetComputeBufferParam(cullingComputeShader, cullTerrainKernel, "_TotalVisibleInstancesOnlyPosWSIDBuffer", totalvisibleInstancesOnlyPosWSIDBuffer);
                cmd.SetComputeMatrixParam(cullingComputeShader, "_VPMatrix", vpPrev);
                cmd.DispatchCompute(cullingComputeShader, cullTerrainKernel, Mathf.CeilToInt(allNodeInfo.Count / 64f), 1, 1);
                //cmd.DispatchCompute(cullingComputeShader, cullTerrainKernel, computeArgBuffer, 0);

                cmd.CopyCounterValue(visibleInstancesOnlyPosWSIDBuffer, argsBuffer, 4);
                cmd.DrawMeshInstancedIndirect(instanceMesh, 0, terrainMat, 2, argsBuffer);
            }
            else
            {
                cmd.CopyCounterValue(invisibleInstancesOnlyPosWSIDBuffer, computeArgBuffer, 0);

                computeArgBuffer.GetData(secondPassComputeArgs);
                if (secondPassComputeArgs[0] != 0)
                {
                    secondPassComputeArgs[0] = (uint)Mathf.CeilToInt(secondPassComputeArgs[0] / 64.0f);
                    //computeArgBuffer.SetData(secondPassComputeArgs);
                    visibleInstancesOnlyPosWSIDBuffer.SetCounterValue(0);
                    cmd.SetComputeTextureParam(cullingComputeShader, cullTerrainKernelSecondPass, "_HiZMap", hizRT);
                    cmd.SetComputeVectorParam(cullingComputeShader, "_HizSize", new Vector4(hizRT.width, hizRT.height, 0, 0));

                    cmd.SetComputeBufferParam(cullingComputeShader, cullTerrainKernelSecondPass, "_InputInstancesOnlyPosWSIDBuffer", invisibleInstancesOnlyPosWSIDBuffer);
                    cmd.SetComputeBufferParam(cullingComputeShader, cullTerrainKernelSecondPass, "_VisibleInstancesOnlyPosWSIDBuffer", visibleInstancesOnlyPosWSIDBuffer);
                    cmd.SetComputeBufferParam(cullingComputeShader, cullTerrainKernelSecondPass, "_TotalVisibleInstancesOnlyPosWSIDBuffer", totalvisibleInstancesOnlyPosWSIDBuffer);
                    cmd.SetComputeMatrixParam(cullingComputeShader, "_VPMatrix", cam.projectionMatrix * cam.worldToCameraMatrix);
                    cmd.DispatchCompute(cullingComputeShader, cullTerrainKernel, Mathf.CeilToInt(secondPassComputeArgs[0] / 64f), 1, 1);
                    //cmd.DispatchCompute(cullingComputeShader, cullTerrainKernelSecondPass, computeArgBuffer, 0);

                    cmd.CopyCounterValue(visibleInstancesOnlyPosWSIDBuffer, argsBuffer, 4);
                    cmd.DrawMeshInstancedIndirect(instanceMesh, 0, terrainMat, 2, argsBuffer);
                }
            }
        }

        context.ExecuteCommandBuffer(cmd);
        CommandBufferPool.Release(cmd);
    }
}
